package com.os.opencv.java.chapter12;

import org.opencv.core.*;
import org.opencv.highgui.HighGui;
import org.opencv.imgproc.Imgproc;
import org.opencv.ml.Ml;
import org.opencv.ml.SVM;

import java.util.Random;

public class SVM2 {

    public static void main(String[] args) {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);

        //参数准备
        int nSample = 100;
        float frac = 0.9f;  //现行部分占比90%
        int nLinear = (int) (frac * nSample);
        Mat trainMat = new Mat(2 * nSample, 2, CvType.CV_32F);
        Mat labels = new Mat(2 * nSample, 1, CvType.CV_32S);
        Random r = new Random(100);

        //创建显示测试结果的图像dst
        int width = 600;
        int height = 600;
        Mat dst = Mat.zeros(height, width, CvType.CV_8UC3);

        //设置线性可训练部分 开始

        //随机生成1类数据的x坐标，范围为图像宽度的0~40%
        Mat trainClass = trainMat.rowRange(0, nLinear);
        Mat c = trainClass.colRange(0, 1);
        float[] cData = new float[90];
        int len = cData.length;
        double[] cDb1 = r.doubles(len, 0, 0.4f*width).toArray();

        for(int i=0; i<len; i++){
            cData[i] = (float) cDb1[i];
        }
        c.put(0,0,cData);

        //随机生成第一类数据的y坐标，范围为图像高度的0~100%
        c = trainClass.colRange(1,2);
        cData = new float[90];
        len = cData.length;
        cDb1 = r.doubles(len, 0, height).toArray();
        for(int i=0; i<len; i++){
            cData[i] = (float) cDb1[i];
        }
        c.put(0,0,cData);

        //随机生成第二类数据的x坐标，范围为图像宽度的60%~100%
        trainClass = trainMat.rowRange(2*nSample-nLinear, 2*nSample);

        c = trainClass.colRange(0, 1);
        cData = new float[90];
        len = cData.length;
        cDb1 = r.doubles(len, 0.6*width, width).toArray();
        for(int i=0; i<len; i++){
            cData[i] = (float) cDb1[i];
        }
        c.put(0,0,cData);

        //随机生成第二类数据的y坐标，范围为图像高度的0~100%；
        c = trainClass.colRange(1, 2);
        cData = new float[90];
        len = cData.length;
        cDb1 = r.doubles(len, 0, height).toArray();
        for(int i=0; i<len; i++){
            cData[i] = (float) cDb1[i];
        }
        c.put(0,0,cData);

        //设置线性可分的训练数据 结束

        //设置线性不可分的训练数据 开始
        //随机生成第一类和第二类数据的坐标点
        trainClass = trainMat.rowRange(nLinear, 2*nSample - nLinear);

        //x坐标范围为图像宽度的40%~60%
        c = trainClass.colRange(0,1);
        cData = new float[20];
        len = cData.length;
        cDb1 = r.doubles(len, 0.4 * width, 0.6 * width).toArray();
        for(int i=0; i<len; i++){
            cData[i] = (float) cDb1[i];
        }
        c.put(0,0,cData);

        //y坐标范围为图像高度0~100%
        c = trainClass.colRange(1, 2);
        cData = new float[20];
        cDb1 = r.doubles(len, 0, height).toArray();
        for(int i=0; i<len; i++){
            cData[i] = (float) cDb1[i];
        }
        c.put(0,0,cData);

        //设置线性不可分的训练数据 结束

        //训练数据标签
        labels.rowRange(0, nSample).setTo(new Scalar(1));
        labels.rowRange(nSample, 2*nSample).setTo(new Scalar(2));

        //创建svm对象并设置参数
        SVM svm = SVM.create();
        svm.setType(SVM.C_SVC);
        svm.setC(0.1);
        svm.setKernel(SVM.LINEAR);
        svm.setTermCriteria(new TermCriteria(TermCriteria.MAX_ITER, (int)1e7, 1e6));

        //训练svm分类器
        svm.train(trainMat, Ml.ROW_SAMPLE, labels);

        //测试每像素并将结果放入结果图像
        int col = dst.cols();
        int cnl = dst.channels();
        byte[] dstData = new byte[(int) (dst.total() * cnl)];
        Mat sampleMat = new Mat(1, 2, CvType.CV_32F);
        float[] sampleData = new float[20];
        for(int i=0; i<dst.rows(); i++){
            for(int j=0; j<col; j++){
                //对每个像素值进行测试
                sampleData[0] = j;
                sampleData[1] = i;
                sampleMat.put(0,0,sampleData);
                float result = svm.predict(sampleMat);

                if(result == 1){
                    dstData[(i*col+j)*cnl] = 0;
                    dstData[(i*col+j)*cnl+1] = 127;
                    dstData[(i*col+j)*cnl+2] = 0;
                }else if(result == 2){
                    dstData[(i*col+j)*cnl] = 127;
                    dstData[(i*col+j)*cnl+1] = 0;
                    dstData[(i*col+j)*cnl+2] = 0;
                }
            }
        }

        //将所有数据放入结果图像中
        dst.put(0,0,dstData);

        //用于生成测试结果图像
        float px, py;
        col = trainMat.cols();

        //第一类用绿色
        float[] trainData = new float[400];
        trainMat.get(0,0,trainData);
        for(int i=0; i<nSample; i++){
            px = trainData[i*col];
            py = trainData[i*col+1];
            Imgproc.circle(dst, new Point(px, py), 3, new Scalar(0,255,0), -1);
        }

        //第二类用蓝色
        for(int i=nSample; i<2*nSample; ++i){
            px = trainData[i*col];
            py = trainData[i*col + 1];
            Imgproc.circle(dst, new Point(px, py), 3, new Scalar(255, 0, 0), -1);
        }

        //获取原始的支持向量
        Mat sv = svm.getUncompressedSupportVectors();
        float[] svData = new float[(int)(sv.total() * sv.channels())];
        sv.get(0,0,svData);

        //在结果图像上用矩形画出支持向量
        for(int i=0; i<sv.rows(); i++){
            double x = svData[i*sv.cols()];
            double y = svData[i*sv.cols()+1];
            Imgproc.rectangle(dst, new Point(x-8, y-8), new Point(x+8, y+8), new Scalar(180, 180, 180), 2);
        }

        //在屏幕上显示结果图像
        HighGui.imshow("svm", dst);
        HighGui.waitKey(0);
        System.exit(0);
    }
}
